#!/usr/bin/env python3
import sys
import argparse
import os
import logging
import wlutil
import contextlib
import shutil
import collections
import pathlib

if not shutil.which('riscv64-unknown-linux-gnu-gcc'):
    sys.exit("No riscv toolchain detected. Please install riscv-tools.")

# Delete a file but don't throw an exception if it doesn't exist
def deleteSafe(pth):
    with contextlib.suppress(FileNotFoundError):
        os.remove(pth)

def main():
    parser = argparse.ArgumentParser(
        description="Build and run (in spike or qemu) boot code and disk images for firesim")
    parser.add_argument('--workdir', help='Use a custom workload directory (defaults to the same directory as the first config file)', type=pathlib.Path)
    parser.add_argument('-v', '--verbose',
                        help='Print all output of subcommands to stdout as well as the logs', action='store_true')
    parser.add_argument('-i', '--initramfs', action='store_true', help="Alias for --no-disk")
    parser.add_argument('-d', '--no-disk', action='store_true', help="Use the no-disk version of this workload (the rootfs will be included as an initramfs in the boot binary)")
    subparsers = parser.add_subparsers(title='Commands', dest='command')

    # Build command
    build_parser = subparsers.add_parser(
        'build', help='Build an image from the given configuration.')
    build_parser.add_argument('config_files', metavar="config", type=pathlib.Path, nargs='+', help="Configuration file(s) to use.")
    build_parser.add_argument('-B', '--binOnly', action='store_true', help="Only build the binary")
    build_parser.add_argument('-I', '--imgOnly', action='store_true', help="Only build the image (may require an image if you have guest-init scripts)")

    # Launch command
    launch_parser = subparsers.add_parser(
        'launch', help='Launch an image on a software simulator (defaults to qemu)')
    launch_parser.add_argument('-s', '--spike', action='store_true',
            help="Use the spike isa simulator instead of qemu")
    launch_parser.add_argument('-j', '--job', nargs='?', default='all',
            help="Launch the specified job. Defaults to running the base image.")
    # the type= option here allows us to only accept one argument but store it
    # in a list so it matches the "build" behavior
    launch_parser.add_argument('config_files', metavar='config', nargs='?', type=(lambda c: [ pathlib.Path(c) ]), help="Configuration file to use.")

    # Test command
    test_parser = subparsers.add_parser(
            'test', help="Test each workload.")
    test_parser.add_argument('config_files', metavar="config", type=pathlib.Path, nargs='+', help="Configuration file(s) to use.")
    test_parser.add_argument('-s', '--spike', action='store_true',
            help="Use the spike isa simulator instead of qemu")
    test_parser.add_argument('-m', '--manual', metavar='testDir', help="Manual test, don't build or run, just compare testDir against the reference output.")

    # Clean Command
    clean_parser = subparsers.add_parser(
            'clean', help="Removes build outputs of the provided config (img and bin). Does not affect logs or runOutputs.")
    clean_parser.add_argument('config_files', metavar="config", type=pathlib.Path, nargs='+', help="Configuration file(s) to use.")

    # Install Command
    install_parser = subparsers.add_parser(
            'install', help="Install this workload to firesim (create configs in firesim/deploy/workloads)")
    install_parser.add_argument('config_files', metavar="config", type=pathlib.Path, nargs='+', help="Configuration file(s) to use.")

    args = parser.parse_args()
    
    # Perform any basic setup functions for wlutil.
    try:
        wlutil.initialize()
    except wlutil.ConfigurationError as e:
        print("Failed to initialize FireMarshal:")
        print(e)
        sys.exit(1)

    ctx = wlutil.getCtx()

    wlutil.initLogging(args.verbose)

    # Load all the configs from the workload directories
    # Order matters here, duplicate workload files found in later search paths
    # will overwrite files found in earlier search paths.

    # We use the keys of an ordered dict to enforce order and uniqueness on the
    # search paths. The value is meaningless.
    workdirs = collections.OrderedDict()

    # board builtin workloads (the *-base.json's).
    workdirs[wlutil.getOpt('workdir-builtin')] = None

    # Configured search paths
    for d in wlutil.getOpt('workload-dirs'):
        workdirs[d] = None

    # User-provided workload search path
    if args.workdir is not None:
        workdirs[args.workdir] = None

    # Treat the workload name as a path and search locally for it
    args.config_files = [ f.resolve() for f in args.config_files ] 
    for f in args.config_files:
        workdirs[f.parent] = None

    cfgs = wlutil.ConfigManager(workdirs.keys())

    if args.command == 'test':
        suitePass = True
        
    skipCount = 0
    failCount = 0
    for cfgPath in args.config_files:
        # Each config gets it's own logging output and results directory
        ctx.setRunName(cfgPath, args.command)
        wlutil.initLogging(args.verbose)
        cfgName = cfgPath.name

        log = logging.getLogger()
        if not args.verbose:
            print("To check on progress, either call marshal with '-v' or see the live output at: ")
            print(wlutil.getOpt('log-dir') / (wlutil.getOpt('run-name') + ".log"))

        try:
            targetCfg = cfgs[cfgName]
        except KeyError as e:
            log.error("Cannot locate workload: " + cfgName)
            sys.exit(1)
   
        if args.initramfs or args.no_disk:
            targetCfg['nodisk'] = True
            if 'jobs' in targetCfg:
                for j in targetCfg['jobs'].values():
                    j['nodisk'] = True

        if args.command == "build":
            if args.binOnly or args.imgOnly:
                # It's fine if they pass -IB, it just builds both
                ret = wlutil.buildWorkload(cfgName, cfgs, buildBin=args.binOnly, buildImg=args.imgOnly)
            else:
                ret = wlutil.buildWorkload(cfgName, cfgs)

            if ret != 0:
                log.error("Failed to build workload " + cfgName)

        elif args.command == "launch":
            # job-configs are named special internally
            if args.job != 'all':
                if 'jobs' in targetCfg: 
                    args.job = targetCfg['name'] + '-' + args.job
                else:
                    log.error("Job " + args.job + " requested, but no jobs specified in config file\n")
                    parser.print_help()

            try:
                outputPath = wlutil.launchWorkload(targetCfg, args.job, args.spike)
            except Exception as e:
                log.exception("Failed to launch workload:")
                outputPath = None

            if outputPath is not None:
                log.info("Workload outputs available at: " + str(outputPath))

        elif args.command == "test":
            log.info("Testing: " + str(cfgPath))
            res, resPath = wlutil.testWorkload(cfgName, cfgs, args.verbose, spike=args.spike, cmp_only=args.manual)
            if res is wlutil.testResult.failure:
                log.info("Test Failed")
                log.info("Output available at: " + str(resPath))
                suitePass = False
                failCount += 1
            elif res is wlutil.testResult.skip:
                log.info("Test Skipped")
                skipCount += 1
            else:
                log.info("Test Passed")
                log.info("Output available at: " + str(resPath))
            log.info("\n")

        elif args.command == 'clean':
            def cleanCfg(cfg):
                if 'bin' in cfg:
                    deleteSafe(cfg['bin'])
                    deleteSafe(wlutil.noDiskPath(cfg['bin']))
                if 'dwarf' in cfg:
                    deleteSafe(cfg['dwarf'])
                    deleteSafe(wlutil.noDiskPath(cfg['dwarf']))
                if 'img' in cfg:
                    deleteSafe(cfg['img'])

            cleanCfg(targetCfg)
            if 'jobs' in targetCfg:
                for jCfg in targetCfg['jobs'].values():
                    cleanCfg(jCfg)

        elif args.command == 'install':
            try:
                wlutil.installWorkload(cfgName, cfgs)
            except wlutil.ConfigurationError as e:
                print(e)
                sys.exit(1)
        else:
            log.error("No subcommand specified")
            sys.exit(1)

    log.info("Log available at: " + str(wlutil.getOpt('log-dir') / (wlutil.getOpt('run-name') + ".log")))
    if args.command == 'test':
        if suitePass:
            log.info("SUCCESS: All Tests Passed (" + str(skipCount) + " tests skipped)")
            sys.exit(0)
        else:
            log.error("FAILURE: " + str(failCount) + " tests failed")
            sys.exit(1)

    sys.exit(0)

if __name__ == "__main__":
    main()
